import os

# 语料库corpus, 二维列表;
# corpus[i][j]: 第i篇文档的第j个单词
corpus = []
vocabulary = {}     # 词汇表,语料库里出现的全部单词, {(string)word: (int)id}
docs_name = []      # 按照语料库里文档的顺序, 将对应的文档名存入该列表
# 读取txt文件来获取单个文档, 并以单词列表的形式返回
def read_txt_get_document(file_path):
	d = []
	# 数据集有些问题, 我理想中的数据集应该是每个文档只有一行,
	# 然后所有单词之间以空格隔开;
	# 但是这个不是, 这个分好多行, 有时候一行的末尾还会有空格, 我也是日了;
	with open(file_path, 'r', encoding='utf-8') as reader:
		for line in reader.readlines():
			if line != '\n':        # 清掉换行符, 后面还清掉了空字符
				for x in line.strip().split(' '):
					if x != '':
						d.append(x)
						word_to_vocabulary(x)   # 填充词汇表
	return d

# 加载路径下的全部txt文档, 并把所有的单词列表写入corpus;
def load_documents(folder_path):
	file_list = os.listdir(folder_path)
	for name in file_list:
		file = os.path.join(folder_path, name)
		if os.path.isfile(file):
			if name.split('.')[1] == "txt":
				# 将文件里的内容加载到列表中
				corpus.append(read_txt_get_document(file))
				docs_name.append(name)


# 检查词汇表, 如果没有当前word, 则新增, 并且将当前词汇表长度作为value
def word_to_vocabulary(word):
	if vocabulary.get(word) == None:
		vocabulary[word] = len(vocabulary)
	return vocabulary[word]

# 把语料库二维数组从String转换成int
def translate_word_to_id():
	for i in range(0, len(corpus)):
		for j in range(0, len(corpus[i])):
			# corpus[i][j] = word_to_vocabulary(corpus[i][j])
			id = vocabulary.get(corpus[i][j])
			if id != None:
				corpus[i][j] = id
			else:       # 新增当前word并返回id
				corpus[i][j] = word_to_vocabulary(corpus[i][j])

if __name__ == "__main__":
	load_documents("../data/test")
	print(corpus)
	print(len(corpus))
	print(vocabulary)
	print(len(vocabulary))
	translate_word_to_id()
	print(corpus)